package org.finesys.chat.core.base.embedding.store.pgvector.handler;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import org.finesys.chat.core.base.document.Metadata;
import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.store.pgvector.MetadataColumnDefinition;
import org.finesys.chat.core.base.embedding.store.pgvector.MetadataHandler;
import org.finesys.chat.core.base.embedding.store.pgvector.MetadataStorageConfig;
import org.finesys.chat.core.base.embedding.store.pgvector.filter.JSONFilterMapper;
import org.finesys.common.core.util.Utils;
import org.finesys.common.core.util.ValidationUtil;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;

public class JSONMetadataHandler implements MetadataHandler {

    private static final ObjectMapper OBJECT_MAPPER;
    final MetadataColumnDefinition columnDefinition;
    final String columnName;
    final JSONFilterMapper filterMapper;
    final List<String> indexes;

    static {
        OBJECT_MAPPER = new ObjectMapper().enable(SerializationFeature.INDENT_OUTPUT);
    }

    public JSONMetadataHandler(MetadataStorageConfig config) {
        List<String> definition = ValidationUtil.ensureNotEmpty(config.columnDefinitions(), "Metadata definition");
        if (definition.size() > 1) {
            throw new IllegalArgumentException("Only one metadata column is supported");
        } else {
            this.columnDefinition = MetadataColumnDefinition.from(definition.get(0));
            this.columnName = this.columnDefinition.getName();
            this.filterMapper = new JSONFilterMapper(columnName);
            this.indexes =Utils.getOrDefault(config.indexes(), Collections.EMPTY_LIST);
        }

    }

    @Override
    public String columnDefinitionsString() {
        return this.columnDefinition.getFullDefinition();
    }

    @Override
    public void createMetadataIndexes(Statement statement, String index) {
        if (!indexes.isEmpty()) {
            throw new RuntimeException("Indexes are not allowed for JSON metadata, use JSONB instead");
        }
    }

    @Override
    public List<String> columnsNames() {
        return Collections.singletonList(this.columnName);
    }

    @Override
    public String whereClause(Filter expression) {
        return this.filterMapper.map(expression);
    }

    @Override
    public Metadata fromResultSet(ResultSet resultSet) {

        try {
            String metadataJson = Utils.getOrDefault(resultSet.getString(this.columnName), "{}");
            return new Metadata(OBJECT_MAPPER.readValue(metadataJson, Map.class));
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public String insertClause() {
        return String.format("%s=EXCLUDED.%s", this.columnName, this.columnName);
    }

    @Override
    public void setMetadata(PreparedStatement preparedStatement, Integer index, Metadata metadata) {
        try {
            preparedStatement.setObject(index, OBJECT_MAPPER.writeValueAsString(metadata.asMap()));
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
    }
}
